package org.kiwi.data.sharding;

import org.springframework.beans.factory.config.BeanDefinitionHolder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.context.annotation.ClassPathBeanDefinitionScanner;
import org.springframework.core.type.classreading.MetadataReader;
import org.springframework.core.type.classreading.MetadataReaderFactory;
import org.springframework.core.type.filter.AnnotationTypeFilter;
import org.springframework.core.type.filter.TypeFilter;

import java.io.IOException;
import java.lang.annotation.Annotation;
import java.util.Set;

/**
 * Created by jack on 16/12/18.
 */
public class ClassPathShardingStrategyScanner extends ClassPathBeanDefinitionScanner {

    private Class<? extends Annotation> shardingStrategyClass;

    public void setShardingStrategyClass(Class<? extends Annotation> shardingStrategyClass) {
        this.shardingStrategyClass = shardingStrategyClass;
    }

    public Class<? extends Annotation> getShardingStrategyClass() {
        return this.shardingStrategyClass;
    }

    public ClassPathShardingStrategyScanner(BeanDefinitionRegistry registry) {
        super(registry, false);
    }

    public Set<BeanDefinitionHolder> doScan(String... basePackages) {
        return super.doScan(basePackages);
    }

    public void registerFilters() {
        boolean acceptAllInterfaces = true;

        if (this.shardingStrategyClass != null) {
            addIncludeFilter(new AnnotationTypeFilter(this.shardingStrategyClass));
            acceptAllInterfaces = false;
        }

        if (acceptAllInterfaces) {
            addIncludeFilter(new TypeFilter() {
                public boolean match(MetadataReader metadataReader, MetadataReaderFactory metadataReaderFactory)
                        throws IOException {
                    return true;
                }
            });
        }

        addExcludeFilter(new TypeFilter() {
            public boolean match(MetadataReader metadataReader, MetadataReaderFactory metadataReaderFactory)
                    throws IOException {
                String className = metadataReader.getClassMetadata().getClassName();
                return className.endsWith("package-info");
            }
        });
    }
}
